import sys
import os
from dotenv import load_dotenv
from openai import OpenAI
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QTextEdit, QPushButton
from PyQt5.QtCore import Qt, QThread, pyqtSignal

# 加载环境变量
load_dotenv()

# AIWorker类重写
class AIWorker(QThread):
    update_signal = pyqtSignal(str)

    def __init__(self, user_input):
        super().__init__()
        self.user_input = user_input
        
        # 初始化 OpenAI 客户端
        self.client = OpenAI(
            base_url=os.getenv('OPENAI_API_BASE'),
            api_key=os.getenv('OPENAI_API_KEY')
        )

    def run(self):
        try:
            # 使用 OpenAI 兼容接口进行对话
            stream = self.client.chat.completions.create(
                model=os.getenv('OPENAI_API_MODEL', 'deepseek-chat'),
                messages=[{'role': 'user', 'content': self.user_input}],
                stream=True
            )

            # 处理流式响应
            for chunk in stream:
                if chunk.choices[0].delta.content is not None:
                    content = chunk.choices[0].delta.content
                    self.update_signal.emit(content)
        except Exception as e:
            self.update_signal.emit(f"\n[Error] {str(e)}")


# 主聊天窗口类
class ChatWindow(QWidget):
    def __init__(self):
        super().__init__()
        # 初始化UI界面
        self.init_ui()
        self.ai_thread = None

    def init_ui(self):
        self.setWindowTitle("InChat")  # 设置窗口标题
        self.setGeometry(100, 100, 600, 500)  # 设置窗口初始大小与位置

        layout = QVBoxLayout()  # 创建垂直布局

        # 创建聊天显示框，设置为只读，方便显示聊天内容
        self.chat_display = QTextEdit()
        self.chat_display.setReadOnly(True)
        self.chat_display.setStyleSheet("""
            background-color: #f0f0f0;
            font-size: 14px;
            padding: 10px;
        """)

        # 创建输入框，用户输入消息
        self.input_area = QTextEdit()
        self.input_area.setMaximumHeight(100)
        self.input_area.setPlaceholderText("输入您的问题...")
        self.input_area.setStyleSheet("""
            background-color: white;
            font-size: 14px;
            padding: 10px;
        """)

        # 创建发送按钮，点击时发送消息
        self.send_btn = QPushButton("提问")
        self.send_btn.setStyleSheet("""
            QPushButton {
                background-color: #4CAF50;
                color: white;
                border: none;
                padding: 10px;
                font-size: 14px;
            }
            QPushButton:hover { background-color: #45a049; }
        """)
        self.send_btn.clicked.connect(self.send_message)  # 点击按钮时调用send_message

        # 将组件添加到布局中
        layout.addWidget(self.chat_display)
        layout.addWidget(self.input_area)
        layout.addWidget(self.send_btn)
        self.setLayout(layout)

    def send_message(self):
        # 获取用户输入的文本并去掉空白
        user_input = self.input_area.toPlainText().strip()
        if not user_input:
            return  # 如果输入为空，返回

        # 在聊天窗口中显示用户输入的消息
        self._append_message("You", user_input)
        self.input_area.clear()  # 清空输入框

        # 在聊天窗口中显示“思考中...”提示
        self._append_message("Bot", "思考中...", is_streaming=True)

        # 启动后台线程，处理与AI的对话
        self.ai_thread = AIWorker(user_input)
        # 连接信号和槽函数，更新聊天内容
        self.ai_thread.update_signal.connect(self.update_bot_response)
        # 连接线程完成后的槽函数
        self.ai_thread.finished.connect(self.finalize_response)
        # 启动线程
        self.ai_thread.start()

    def _append_message(self, sender, message, is_streaming=False):
        # 获取文本光标，确保内容插入到聊天窗口的末尾
        cursor = self.chat_display.textCursor()
        cursor.movePosition(cursor.End)

        # 设置消息的前缀和颜色，用户消息和机器人消息区分开
        if sender == "You":
            prefix = "\nYou: "
            color = "#333333"  # 用户消息颜色
        else:
            prefix = "\nBot: "
            color = "#2196F3" if not is_streaming else "#888888"  # 机器人消息颜色，正在处理时颜色较灰

        # 插入消息
        cursor.insertText(prefix)
        cursor.insertHtml(f'<span style="color:{color}">{message}</span>')
        # 确保文本框显示最新内容
        self.chat_display.ensureCursorVisible()

    def update_bot_response(self, content):
        # 如果聊天框正在显示“思考中...”提示，则替换为实际响应内容
        current_text = self.chat_display.toPlainText()
        if current_text.endswith("思考中..."):
            self.chat_display.setPlainText(current_text[:-7] + content)
        else:
            # 否则直接插入新的消息
            cursor = self.chat_display.textCursor()
            cursor.movePosition(cursor.End)  # 确保插入文本到末尾
            cursor.insertText(content)  # 插入文本
        self.chat_display.ensureCursorVisible()

    def finalize_response(self):
        # 确保最后有换行符，使聊天内容整洁
        current = self.chat_display.toPlainText()
        if not current.endswith("\n"):
            self.chat_display.append("")


# 主程序入口
if __name__ == "__main__":
    app = QApplication(sys.argv)  # 创建应用程序
    window = ChatWindow()  # 创建聊天窗口实例
    window.show()  # 显示窗口
    sys.exit(app.exec_())  # 进入应用的主循环

